import os
from PIL import Image
import pytesseract
from pdf2image import convert_from_path
import argparse
import sys
import numpy as np

# 尝试导入PaddleOCR
try:
    from paddleocr import PaddleOCR
    PADDLE_AVAILABLE = True
except ImportError:
    PADDLE_AVAILABLE = False
    print("警告: PaddleOCR未安装，将使用Tesseract进行OCR", file=sys.stderr)

def preprocess_image(image):
    """
    预处理图片以提高OCR准确率
    
    Args:
        image: PIL Image对象
    
    Returns:
        处理后的PIL Image对象
    """
    # 转换为灰度图
    image = image.convert('L')
    
    # 增加对比度
    from PIL import ImageEnhance
    enhancer = ImageEnhance.Contrast(image)
    image = enhancer.enhance(2.0)
    
    return image

def images_to_pdf(input_files, output_path, dpi=200):
    """
    将多个图片合并成一个PDF文件
    
    Args:
        input_files (list): 输入图片文件路径列表
        output_path (str): 输出PDF文件路径
        dpi (int): 图片DPI
    """
    try:
        # 确保输出目录存在
        output_dir = os.path.dirname(output_path)
        if output_dir:
            os.makedirs(output_dir, exist_ok=True)
        
        # 检查输入文件是否存在
        for img_path in input_files:
            if not os.path.exists(img_path):
                raise FileNotFoundError(f"找不到输入文件: {img_path}")
        
        # 打开第一张图片
        first_image = Image.open(input_files[0]).convert('RGB')
        images = []
        
        # 转换其他图片
        for img_path in input_files[1:]:
            img = Image.open(img_path).convert('RGB')
            images.append(img)
        
        # 保存为PDF
        first_image.save(output_path, save_all=True, append_images=images, dpi=(dpi, dpi))
        print(f'已创建PDF文件: {output_path}')
        
    except Exception as e:
        print(f"错误: {str(e)}", file=sys.stderr)
        sys.exit(1)

def image_to_text(input_file, output_path, lang='ch', use_paddle=True):
    """
    将图片转换为文字
    
    Args:
        input_file (str): 输入图片文件路径
        output_path (str): 输出文本文件路径
        lang (str): OCR语言，默认中文
        use_paddle (bool): 是否使用PaddleOCR
    """
    try:
        # 确保输出目录存在
        output_dir = os.path.dirname(output_path)
        if output_dir:
            os.makedirs(output_dir, exist_ok=True)
        
        # 检查输入文件是否存在
        if not os.path.exists(input_file):
            raise FileNotFoundError(f"找不到输入文件: {input_file}")
        
        # 打开图片
        image = Image.open(input_file)
        
        # 如果PaddleOCR不可用，强制使用Tesseract
        if use_paddle and not PADDLE_AVAILABLE:
            print("警告: PaddleOCR未安装，将使用Tesseract进行OCR", file=sys.stderr)
            use_paddle = False
        
        if use_paddle:
            # 使用PaddleOCR
            ocr = PaddleOCR(use_angle_cls=True, lang=lang, show_log=False)
            result = ocr.ocr(np.array(image), cls=True)
            
            # 提取文本
            text = ""
            if result:
                for line in result:
                    for item in line:
                        text += item[1][0] + "\n"
        else:
            # 使用Tesseract
            # 预处理图片
            processed_image = preprocess_image(image)
            
            # 使用pytesseract进行OCR
            text = pytesseract.image_to_string(processed_image, lang=lang)
            
            # 如果文本为空，尝试使用原始图片
            if not text.strip():
                text = pytesseract.image_to_string(image, lang=lang)
        
        # 保存文本
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write(text)
        
        print(f'已创建文本文件: {output_path}')
        
        # 如果文本为空，显示警告
        if not text.strip():
            print("警告: 未能识别出任何文字，请检查图片质量", file=sys.stderr)
        
    except Exception as e:
        print(f"错误: {str(e)}", file=sys.stderr)
        sys.exit(1)

def main():
    parser = argparse.ArgumentParser(description='图片处理工具')
    subparsers = parser.add_subparsers(dest='command', help='子命令')
    
    # 图片转PDF命令
    pdf_parser = subparsers.add_parser('to-pdf', help='将图片转换为PDF')
    pdf_parser.add_argument('input_files', nargs='+', help='输入图片文件路径（可以指定多个）')
    pdf_parser.add_argument('-o', '--output', default='output.pdf', help='输出PDF文件路径（默认：output.pdf）')
    pdf_parser.add_argument('--dpi', type=int, default=200, help='图片DPI（默认：200）')
    
    # 图片转文字命令
    text_parser = subparsers.add_parser('to-text', help='将图片转换为文字')
    text_parser.add_argument('input_file', help='输入图片文件路径')
    text_parser.add_argument('-o', '--output', default='output.txt', help='输出文本文件路径（默认：output.txt）')
    text_parser.add_argument('--lang', default='ch', help='OCR语言（默认：ch，可选：ch, en, french, german, korean, japan）')
    text_parser.add_argument('--use-tesseract', action='store_true', help='使用Tesseract而不是PaddleOCR')
    
    args = parser.parse_args()
    
    if args.command == 'to-pdf':
        images_to_pdf(args.input_files, args.output, args.dpi)
    elif args.command == 'to-text':
        image_to_text(args.input_file, args.output, args.lang, not args.use_tesseract)
    else:
        parser.print_help()

if __name__ == '__main__':
    main() 